#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import os
import json
import shutil
from PyQt5.QtWidgets import (QApplication, QMainWindow, QTabWidget, QVBoxLayout, 
                            QWidget, QLabel, QStatusBar, QAction, QFileDialog,
                            QMessageBox, QSplashScreen, QInputDialog, QDialog,
                            QFormLayout, QLineEdit, QPushButton, QHBoxLayout,
                            QListWidget, QTreeWidget, QTreeWidgetItem, QSplitter,
                            QGroupBox, QMenu, QProgressDialog, QFrame)
from PyQt5.QtGui import QIcon, QPixmap, QFont
from PyQt5.QtCore import Qt, QSize, QTimer, QDir, QDateTime, QEvent

# 导入应用程序各模块UI面板
from ui.annotation_panel import AnnotationPanel
from ui.model_manager_panel import ModelManagerPanel
from ui.config_panel import ConfigPanel
from ui.training_panel import TrainingPanel  # 导入训练面板
from ui.inference_panel import InferencePanel  # 导入推理面板

# 检查是否安装了ultralytics库
try:
    import ultralytics
    ULTRALYTICS_AVAILABLE = True
except ImportError:
    ULTRALYTICS_AVAILABLE = False
    print("警告: 未找到ultralytics包，请安装: pip install ultralytics")

class ProjectDialog(QDialog):
    """项目创建对话框"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("创建新项目")
        self.setMinimumWidth(500)
        self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)
        self.init_ui()
        
    def init_ui(self):
        """初始化UI"""
        layout = QVBoxLayout(self)
        layout.setContentsMargins(24, 24, 24, 24)
        layout.setSpacing(20)
        
        # 标题标签
        title_label = QLabel("创建新项目")
        title_label.setProperty("title", "true")
        layout.addWidget(title_label)
        
        # 表单布局
        form_layout = QFormLayout()
        form_layout.setLabelAlignment(Qt.AlignRight)
        form_layout.setSpacing(16)
        
        # 项目名称
        self.name_edit = QLineEdit()
        self.name_edit.setPlaceholderText("输入项目名称...")
        form_layout.addRow("项目名称:", self.name_edit)
        
        # 项目路径
        path_layout = QHBoxLayout()
        path_layout.setSpacing(8)
        self.path_edit = QLineEdit()
        self.path_edit.setReadOnly(True)
        self.path_edit.setPlaceholderText("选择项目存储位置...")
        browse_btn = QPushButton("浏览...")
        browse_btn.setProperty("secondary", "true")
        browse_btn.setFixedWidth(100)
        browse_btn.clicked.connect(self.browse_path)
        path_layout.addWidget(self.path_edit)
        path_layout.addWidget(browse_btn)
        
        form_layout.addRow("项目位置:", path_layout)
        
        # 项目描述
        self.desc_edit = QLineEdit()
        self.desc_edit.setPlaceholderText("输入项目描述...")
        form_layout.addRow("项目描述:", self.desc_edit)
        
        layout.addLayout(form_layout)
        
        # 按钮布局
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(12)
        cancel_btn = QPushButton("取消")
        cancel_btn.setProperty("secondary", "true")
        cancel_btn.setFixedWidth(120)
        cancel_btn.clicked.connect(self.reject)
        
        ok_btn = QPushButton("创建项目")
        ok_btn.setFixedWidth(120)
        ok_btn.clicked.connect(self.accept)
        
        btn_layout.addStretch()
        btn_layout.addWidget(cancel_btn)
        btn_layout.addWidget(ok_btn)
        
        layout.addLayout(btn_layout)
        
    def browse_path(self):
        """浏览项目路径"""
        path = QFileDialog.getExistingDirectory(self, "选择项目位置")
        if path:
            self.path_edit.setText(path)
            
    def get_project_info(self):
        """获取项目信息"""
        return {
            "name": self.name_edit.text(),
            "path": self.path_edit.text(),
            "description": self.desc_edit.text(),
            "created_time": QDateTime.currentDateTime().toString()
        }

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        
        # 设置窗口属性
        self.setWindowTitle("YOLOv10 可视化工具")
        self.setMinimumSize(1200, 800)
        
        # 初始化变量
        self.current_project = None
        self.projects_list = []
        self.max_recent_projects = 5
        
        # 创建状态栏
        self.statusBar = self.statusBar()
        self.statusBar.showMessage("准备就绪")
        
        # 创建项目信息标签（状态栏中）
        self.status_project_label = QLabel("未打开项目")
        self.statusBar.addPermanentWidget(self.status_project_label)
        
        # 创建菜单栏
        self.create_menu_bar()
        
        # 初始化UI
        self.init_ui()
        
        # 加载项目历史
        self.load_projects_history()
        
        # 更新最近项目菜单
        self.update_recent_projects_menu()
        
        # 尝试加载Ultralytics
        try:
            import ultralytics
            self.has_ultralytics = True
        except ImportError:
            self.has_ultralytics = False
            self.statusBar.showMessage("警告: Ultralytics库未安装，某些功能将不可用")
            QTimer.singleShot(3000, self.install_ultralytics)
            
        # 安装事件过滤器，处理快捷键冲突问题
        self.installEventFilter(self)
            
    def eventFilter(self, obj, event):
        """全局事件过滤器，处理快捷键冲突问题"""
        if event.type() == QEvent.KeyPress:
            # 获取当前激活的窗口部件
            focused_widget = QApplication.focusWidget()
            
            # 如果当前焦点在标注面板上，则让标注面板处理快捷键
            if focused_widget and isinstance(focused_widget, QWidget) and hasattr(self, 'annotation_panel'):
                # 如果标注面板被显示且标签页在标注面板上
                if self.tabs.currentWidget() == self.annotation_panel:
                    # 只处理跟标注相关的快捷键
                    if (event.key() in [Qt.Key_Plus, Qt.Key_Equal, Qt.Key_Minus, Qt.Key_0, 
                                       Qt.Key_A, Qt.Key_D] or
                        (event.modifiers() == Qt.ControlModifier and 
                         event.key() in [Qt.Key_S, Qt.Key_Z, Qt.Key_O, Qt.Key_B, Qt.Key_Delete])):
                        # 让事件继续传递到标注面板
                        return False
        
        # 默认处理
        return super().eventFilter(obj, event)
    
    def create_menu_bar(self):
        """创建菜单栏"""
        menu_bar = self.menuBar()
        
        # 文件菜单
        file_menu = menu_bar.addMenu("文件")
        
        new_project_action = QAction(QIcon(), "新建项目", self)
        new_project_action.setShortcut("Ctrl+N")
        new_project_action.triggered.connect(self.create_new_project)
        file_menu.addAction(new_project_action)
        
        open_project_action = QAction(QIcon(), "打开项目", self)
        open_project_action.setShortcut("Ctrl+O")
        open_project_action.triggered.connect(self.open_project)
        file_menu.addAction(open_project_action)
        
        save_project_action = QAction(QIcon(), "保存项目", self)
        save_project_action.setShortcut("Ctrl+S")
        save_project_action.triggered.connect(self.save_project)
        file_menu.addAction(save_project_action)
        
        # 添加导出数据集菜单项
        export_dataset_action = QAction(QIcon(), "导出数据集", self)
        export_dataset_action.setShortcut("Ctrl+E")
        export_dataset_action.triggered.connect(self.export_current_project_dataset)
        file_menu.addAction(export_dataset_action)
        
        file_menu.addSeparator()
        
        # 添加最近项目子菜单
        self.recent_projects_menu = QMenu("最近项目", self)
        file_menu.addMenu(self.recent_projects_menu)
        
        file_menu.addSeparator()
        
        exit_action = QAction(QIcon(), "退出", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)
        
        # 工具菜单
        tools_menu = menu_bar.addMenu("工具")
        
        # 添加训练菜单项
        train_action = QAction(QIcon(), "训练模型", self)
        train_action.setShortcut("Ctrl+T")
        train_action.triggered.connect(self.open_training_panel)
        tools_menu.addAction(train_action)
        
        # 添加推理菜单项
        inference_action = QAction(QIcon(), "模型推理", self)
        inference_action.setShortcut("Ctrl+I")
        inference_action.triggered.connect(self.open_inference_panel)
        tools_menu.addAction(inference_action)
        
        settings_action = QAction(QIcon(), "设置", self)
        settings_action.triggered.connect(self.show_settings)
        tools_menu.addAction(settings_action)
        
        # 帮助菜单
        help_menu = menu_bar.addMenu("帮助")
        
        about_action = QAction(QIcon(), "关于", self)
        about_action.triggered.connect(self.show_about)
        help_menu.addAction(about_action)
        
        # Ultralytics状态菜单项
        if not ULTRALYTICS_AVAILABLE:
            install_ultralytics_action = QAction(QIcon(), "安装Ultralytics", self)
            install_ultralytics_action.triggered.connect(self.install_ultralytics)
            help_menu.addAction(install_ultralytics_action)
        
    def init_ui(self):
        """初始化主UI组件"""
        # 创建中央窗口部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(16, 16, 16, 16)
        main_layout.setSpacing(20)  # 增加整体间距
        
        # 添加项目信息标签
        self.project_info_label = QLabel("未打开项目")
        self.project_info_label.setProperty("title", "true")
        self.project_info_label.setObjectName("project_info_label")
        main_layout.addWidget(self.project_info_label)
        
        # # 添加间隔 - 在项目标签和标签页之间添加间距
        # spacer = QWidget()
        # spacer.setFixedHeight(16)  # 16像素的间隔
        # main_layout.addWidget(spacer)
        
        # 创建标签页控件
        self.tabs = QTabWidget()
        self.tabs.setDocumentMode(True)
        self.tabs.setTabPosition(QTabWidget.North)
        
        # 创建各功能面板
        self.model_manager_panel = ModelManagerPanel()
        self.annotation_panel = AnnotationPanel()
        self.config_panel = ConfigPanel()
        self.training_panel = TrainingPanel()  # 添加训练面板
        self.inference_panel = InferencePanel()  # 添加推理面板
        
        # 添加标签页
        self.tabs.addTab(self.model_manager_panel, "模型管理")
        self.tabs.addTab(self.annotation_panel, "数据标注")
        self.tabs.addTab(self.config_panel, "模型配置")
        self.tabs.addTab(self.training_panel, "训练")  # 添加训练标签页
        self.tabs.addTab(self.inference_panel, "推理")  # 添加推理标签页
        
        # 将标签页添加到主布局
        main_layout.addWidget(self.tabs)
        
        # 连接标签页切换信号
        self.tabs.currentChanged.connect(self.tab_changed)
        
        # 连接模型下载信号
        self.model_manager_panel.model_downloaded.connect(self.on_model_downloaded)
    
    def create_new_project(self):
        """创建新项目"""
        dialog = ProjectDialog(self)
        if dialog.exec_() == QDialog.Accepted:
            project_info = dialog.get_project_info()
            
            # 验证项目信息
            if not project_info["name"]:
                QMessageBox.warning(self, "警告", "项目名称不能为空")
                return
                
            if not project_info["path"]:
                QMessageBox.warning(self, "警告", "项目路径不能为空")
                return
                
            # 检查同一目录下是否已存在同名项目
            project_dir = os.path.join(project_info["path"], project_info["name"])
            if os.path.exists(project_dir):
                config_path = os.path.join(project_dir, "project.json")
                if os.path.exists(config_path):
                    QMessageBox.warning(self, "警告", f"项目 '{project_info['name']}' 已存在于选择的目录中")
                    return
                    
            # 创建项目目录结构
            os.makedirs(project_dir, exist_ok=True)
            
            # 创建项目子目录
            os.makedirs(os.path.join(project_dir, "images"), exist_ok=True)
            os.makedirs(os.path.join(project_dir, "labels"), exist_ok=True)
            os.makedirs(os.path.join(project_dir, "configs"), exist_ok=True)
            os.makedirs(os.path.join(project_dir, "dataset"), exist_ok=True)
            os.makedirs(os.path.join(project_dir, "models"), exist_ok=True)
            
            # 保存项目配置文件
            project_config = {
                "name": project_info["name"],
                "description": project_info["description"],
                "path": project_dir,
                "created_time": project_info["created_time"]
            }
            
            config_path = os.path.join(project_dir, "project.json")
            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(project_config, f, indent=4, ensure_ascii=False)
                
            # 设置为当前项目
            self.current_project = project_config
            
            # 更新项目信息标签
            self.update_project_info_label()
            
            # 添加到项目历史（先移除可能存在的同名项目）
            self.projects_list = [p for p in self.projects_list if p["name"] != project_config["name"]]
            self.projects_list.append(project_config)
            self.save_projects_history()
                
            # 更新最近项目菜单
            self.update_recent_projects_menu()
            
            self.statusBar.showMessage(f"项目 '{project_info['name']}' 创建成功")
            
    def update_project_info_label(self):
        """更新项目信息标签"""
        if self.current_project:
            project_name = f"当前项目: {self.current_project['name']}"
            self.project_info_label.setText(project_name)
            self.status_project_label.setText(project_name)
        else:
            self.project_info_label.setText("未打开项目")
            self.status_project_label.setText("未打开项目")
            
    def open_project(self):
        """打开已有项目"""
        project_dir = QFileDialog.getExistingDirectory(
            self, "选择项目文件夹", "", 
            QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks
        )
        
        if not project_dir:
            return
            
        # 查找项目配置文件
        config_path = os.path.join(project_dir, "project.json")
        if not os.path.exists(config_path):
            # 尝试在父目录查找
            parent_dir = os.path.dirname(project_dir)
            config_path = os.path.join(parent_dir, "project.json")
            if not os.path.exists(config_path):
                reply = QMessageBox.question(
                    self, 
                    "未找到项目配置", 
                    "所选目录不是有效的项目目录，是否创建新项目？",
                    QMessageBox.Yes | QMessageBox.No
                )
                
                if reply == QMessageBox.Yes:
                    # 创建新项目
                    name, ok = QInputDialog.getText(self, "项目名称", "请输入项目名称:")
                    if ok and name:
                        # 检查同一目录下是否有同名项目
                        new_project_dir = os.path.join(project_dir, name)
                        if os.path.exists(new_project_dir):
                            check_config_path = os.path.join(new_project_dir, "project.json")
                            if os.path.exists(check_config_path):
                                QMessageBox.warning(self, "警告", f"项目 '{name}' 已存在于选择的目录中")
                                return
                                
                        project_info = {
                            "name": name,
                            "path": project_dir,
                            "description": "",
                            "created_time": QDateTime.currentDateTime().toString()
                        }
                        
                        # 创建项目结构
                        os.makedirs(os.path.join(project_dir, "images"), exist_ok=True)
                        os.makedirs(os.path.join(project_dir, "labels"), exist_ok=True)
                        os.makedirs(os.path.join(project_dir, "models"), exist_ok=True)
                        os.makedirs(os.path.join(project_dir, "configs"), exist_ok=True)
                        os.makedirs(os.path.join(project_dir, "dataset"), exist_ok=True)
                        
                        # 保存项目配置
                        with open(os.path.join(project_dir, "project.json"), 'w', encoding='utf-8') as f:
                            json.dump(project_info, f, indent=4, ensure_ascii=False)
                            
                        # 设置为当前项目
                        self.current_project = project_info
                        
                        # 更新项目信息标签
                        self.update_project_info_label()
                        
                        # 添加到项目历史（删除同名项目）
                        self.projects_list = [p for p in self.projects_list if p["name"] != project_info["name"]]
                        self.projects_list.append(project_info)
                        self.save_projects_history()
                            
                        # 更新最近项目菜单
                        self.update_recent_projects_menu()
                        
                        self.statusBar.showMessage(f"项目 '{name}' 创建成功")
                return
                
        try:
            # 加载项目配置
            with open(config_path, 'r', encoding='utf-8') as f:
                project_info = json.load(f)
                
            # 设置为当前项目
            self.current_project = project_info
            
            # 更新项目信息标签
            self.update_project_info_label()
            
            # 添加到项目历史（删除同名项目）
            self.projects_list = [p for p in self.projects_list if p["name"] != project_info["name"]]
            self.projects_list.append(project_info)
            self.save_projects_history()
                
            # 更新最近项目菜单
            self.update_recent_projects_menu()
            
            # 更新UI
            self.statusBar.showMessage(f"项目 '{project_info['name']}' 已打开")
            
            # 加载类别定义
            if hasattr(self, 'annotation_panel'):
                self.annotation_panel.load_classes_from_project()
            
            # 检查图像和标签目录，但不自动加载
            images_dir = os.path.join(project_info["path"], "images")
            
            if os.path.exists(images_dir) and os.listdir(images_dir):
                self.statusBar.showMessage(f"项目 '{project_info['name']}' 已打开，含有图像可用于标注")
                
        except Exception as e:
            QMessageBox.critical(self, "错误", f"打开项目时出错: {str(e)}")
                
    def save_project(self):
        """保存当前项目"""
        if not self.current_project:
            QMessageBox.warning(self, "警告", "没有当前项目")
            return
            
        try:
            # 保存标注
            if hasattr(self.annotation_panel, 'canvas') and hasattr(self.annotation_panel.canvas, 'image_path'):
                self.annotation_panel.save_annotations()
                
            # 保存配置
            # TODO: 实现配置保存逻辑
            
            self.statusBar.showMessage(f"项目 '{self.current_project['name']}' 已保存")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存项目时出错: {str(e)}")
        
    def save_projects_history(self):
        """保存项目历史到配置文件"""
        try:
            # 只保存基本信息，不包含完整路径等信息
            history_list = []
            for project in self.projects_list:
                history_list.append({
                    "name": project["name"],
                    "path": project["path"],
                    "description": project.get("description", "")
                })
                
            # 保存到配置文件
            config_dir = os.path.join(os.path.expanduser("~"), ".yolo_vis")
            os.makedirs(config_dir, exist_ok=True)
            
            with open(os.path.join(config_dir, "projects.json"), 'w', encoding='utf-8') as f:
                json.dump(history_list, f, indent=4, ensure_ascii=False)
                
        except Exception as e:
            print(f"保存项目历史失败: {str(e)}")
    
    def load_projects_history(self):
        """加载项目历史记录"""
        try:
            config_path = os.path.join(os.path.expanduser("~"), ".yolo_vis", "projects.json")
            if os.path.exists(config_path):
                with open(config_path, 'r', encoding='utf-8') as f:
                    history_list = json.load(f)
                    
                # 加载项目列表
                for project in history_list:
                    # 检查项目目录是否存在
                    if os.path.exists(project["path"]):
                        self.projects_list.append(project)
                
                # 更新项目信息标签（不自动加载最后一个项目）
                self.update_project_info_label()
        except Exception as e:
            print(f"加载项目历史失败: {str(e)}")
        
    def tab_changed(self, index):
        """标签页切换时的处理函数"""
        tab_name = self.tabs.tabText(index)
        self.statusBar.showMessage(f"当前功能区域: {tab_name}")
        
    def show_settings(self):
        """显示设置对话框"""
        # TODO: 实现设置对话框
        QMessageBox.information(self, "设置", "设置对话框将在此显示")
        
    def show_about(self):
        """显示关于对话框"""
        QMessageBox.about(
            self, "关于 YOLOv10 可视化应用程序",
            "YOLOv10 可视化应用程序\n\n"
            "基于 Ultralytics 的 YOLOv10 可视化应用程序，"
            "实现数据标注、模型配置、训练和推理等功能。\n\n"
            "版本: 1.0.0\n"
            "许可证: MIT"
        )
        
    def install_ultralytics(self):
        """提示安装Ultralytics库"""
        reply = QMessageBox.question(
            self, 
            "安装Ultralytics", 
            "需要安装Ultralytics库才能使用完整功能。\n\n"
            "请在命令行运行:\n"
            "pip install ultralytics\n\n"
            "然后重启应用程序。\n\n"
            "是否打开安装说明？",
            QMessageBox.Yes | QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            # 简单起见，这里只显示一条消息
            QMessageBox.information(
                self,
                "安装说明",
                "请在命令行终端中执行以下命令：\n\n"
                "pip install ultralytics\n\n"
                "安装完成后重启本应用。"
            )
            
    def on_model_downloaded(self, model_name):
        """当模型下载完成时的回调"""
        self.statusBar.showMessage(f"模型 {model_name} 已成功下载")
        
        # 如果有当前项目，将模型复制到项目目录
        if self.current_project:
            try:
                src_model_path = os.path.join(os.getcwd(), "models", f"{model_name.split(' ')[0].lower()}.pt")
                if os.path.exists(src_model_path):
                    project_models_dir = os.path.join(self.current_project["path"], "models")
                    os.makedirs(project_models_dir, exist_ok=True)
                    dst_model_path = os.path.join(project_models_dir, os.path.basename(src_model_path))
                    shutil.copy2(src_model_path, dst_model_path)
            except Exception as e:
                print(f"复制模型到项目目录失败: {str(e)}")

    def update_recent_projects_menu(self):
        """更新最近项目菜单"""
        self.recent_projects_menu.clear()
        
        # 按照项目路径去重
        unique_projects = {}
        for project in self.projects_list:
            # 如果有多个同名项目，仅保留最新的一个
            if project["name"] not in unique_projects:
                unique_projects[project["name"]] = project
                
        # 转换为列表并倒序，最新的项目排在前面
        recent_projects = list(unique_projects.values())
        # 只显示最近的N个项目
        recent_projects = recent_projects[-self.max_recent_projects:]
        recent_projects.reverse()
        
        for project in recent_projects:
            action = QAction(project["name"], self)
            # 使用 lambda 函数时需要设置默认参数捕获当前项目，否则所有操作都会指向最后一个项目
            action.triggered.connect(lambda checked=False, p=project: self.open_project_from_path(p["path"]))
            self.recent_projects_menu.addAction(action)
            
        # 如果没有最近项目，禁用菜单
        self.recent_projects_menu.setEnabled(len(recent_projects) > 0)
    
    def open_project_from_path(self, project_path):
        """从路径打开项目"""
        if os.path.exists(project_path):
            config_path = os.path.join(project_path, "project.json")
            if os.path.exists(config_path):
                try:
                    # 加载项目配置
                    with open(config_path, 'r', encoding='utf-8') as f:
                        project_info = json.load(f)
                        
                    # 设置为当前项目
                    self.current_project = project_info
                    
                    # 更新项目信息标签
                    self.update_project_info_label()
                    
                    # 先删除同名项目，再添加到项目列表
                    self.projects_list = [p for p in self.projects_list if p["name"] != project_info["name"]]
                    self.projects_list.append(project_info)
                    self.save_projects_history()
                    
                    # 更新UI
                    self.statusBar.showMessage(f"项目 '{project_info['name']}' 已打开")
                    
                    # 更新最近项目菜单
                    self.update_recent_projects_menu()
                    
                    # 加载类别定义
                    if hasattr(self, 'annotation_panel'):
                        self.annotation_panel.load_classes_from_project()
                    
                    # 检查图像目录，但不自动加载
                    images_dir = os.path.join(project_path, "images")
                    if os.path.exists(images_dir) and os.listdir(images_dir):
                        self.statusBar.showMessage(f"项目 '{project_info['name']}' 已打开，含有图像可用于标注")
                    
                    return True
                except Exception as e:
                    QMessageBox.critical(self, "错误", f"打开项目时出错: {str(e)}")
            else:
                QMessageBox.warning(self, "警告", f"无效的项目目录: {project_path}")
        else:
            QMessageBox.warning(self, "警告", f"项目路径不存在: {project_path}")
        
        return False
    
    def export_current_project_dataset(self):
        """导出当前项目的数据集"""
        if not self.current_project:
            QMessageBox.warning(self, "警告", "没有打开的项目，请先打开或创建一个项目")
            return
            
        # 获取项目路径
        project_path = self.current_project["path"]
        
        # 检查图像目录
        images_dir = os.path.join(project_path, "images")
        if not os.path.exists(images_dir) or not os.listdir(images_dir):
            QMessageBox.warning(self, "警告", "项目图像目录为空，无法导出数据集")
            return
            
        # 创建项目数据集目录（如果不存在）
        dataset_dir = os.path.join(project_path, "dataset")
        os.makedirs(dataset_dir, exist_ok=True)
        
        # 切换到标注面板并加载图像
        self.tabs.setCurrentIndex(1)  # 切换到标注面板
        # 从项目目录加载图像时不需要复制，直接使用
        self.annotation_panel.open_batch_mode(images_dir)
        
        # 创建数据集结构
        dataset_images_dir = os.path.join(dataset_dir, "images")
        dataset_labels_dir = os.path.join(dataset_dir, "labels")
        
        os.makedirs(dataset_images_dir, exist_ok=True)
        os.makedirs(dataset_labels_dir, exist_ok=True)
        
        # 保存标注
        if hasattr(self.annotation_panel, 'canvas') and hasattr(self.annotation_panel.canvas, 'image_path'):
            self.annotation_panel.save_annotations()
            
        # 创建进度对话框
        progress = QProgressDialog("正在导出数据集...", "取消", 0, 100, self)
        progress.setWindowTitle("导出数据集")
        progress.setWindowModality(Qt.WindowModal)
        progress.setValue(0)
        
        # 批量模式 - 导出所有已加载的图像和标注
        if hasattr(self.annotation_panel, 'image_files') and self.annotation_panel.image_files:
            total_files = len(self.annotation_panel.image_files)
            
            # 显示进度对话框
            progress.setRange(0, total_files)
            
            # 计数器
            exported_images = 0
            exported_labels = 0
            
            # 导出所有图像和对应的标签
            import shutil
            for i, image_path in enumerate(self.annotation_panel.image_files):
                # 检查是否取消
                if progress.wasCanceled():
                    break
                    
                # 更新进度
                progress.setValue(i)
                
                # 复制图像
                image_basename = os.path.basename(image_path)
                image_dest = os.path.join(dataset_images_dir, image_basename)
                try:
                    shutil.copy2(image_path, image_dest)
                    exported_images += 1
                except Exception as e:
                    print(f"复制图像失败: {image_path} - {str(e)}")
                    continue
                    
                # 复制对应的标签文件
                base_name_without_ext = os.path.splitext(image_basename)[0]
                image_dir = os.path.dirname(image_path)
                
                # 确定标签文件路径
                if "images" in image_dir and image_dir.endswith("images"):
                    # 项目结构
                    project_dir = os.path.dirname(image_dir)
                    source_label_path = os.path.join(project_dir, "labels", f"{base_name_without_ext}.txt")
                else:
                    # 常规结构
                    source_label_path = os.path.join(image_dir, "labels", f"{base_name_without_ext}.txt")
                
                # 如果标签文件存在，则复制
                if os.path.exists(source_label_path):
                    label_dest = os.path.join(dataset_labels_dir, f"{base_name_without_ext}.txt")
                    try:
                        shutil.copy2(source_label_path, label_dest)
                        exported_labels += 1
                    except Exception as e:
                        print(f"复制标签失败: {source_label_path} - {str(e)}")
                
            # 完成进度
            progress.setValue(total_files)
            
            # 创建数据集配置文件
            dataset_yaml = os.path.join(dataset_dir, "dataset.yaml")
            class_names = [self.annotation_panel.class_combo.itemText(i) for i in range(self.annotation_panel.class_combo.count())]
            
            with open(dataset_yaml, 'w', encoding='utf-8') as f:
                f.write(f"path: {dataset_dir}\n")
                f.write("train: images\n")
                f.write("val: images\n\n")
                f.write(f"nc: {len(class_names)}\n")
                f.write(f"names: {class_names}\n")
                
            self.statusBar.showMessage(f"数据集已导出: {exported_images}张图像, {exported_labels}个标签文件")
            
            # 如果训练面板已打开，自动加载数据集
            if hasattr(self, 'training_panel'):
                reply = QMessageBox.question(
                    self, 
                    "使用数据集", 
                    "数据集已导出，是否跳转到训练页面使用此数据集？",
                    QMessageBox.Yes | QMessageBox.No
                )
                
                if reply == QMessageBox.Yes:
                    # 跳转到训练面板
                    self.open_training_panel()
                else:
                    # 打开数据集目录
                    os.startfile(dataset_dir)
            else:
                # 打开数据集目录
                os.startfile(dataset_dir)
        else:
            QMessageBox.warning(self, "警告", "没有加载的图像，请先加载图像")
            return

    def open_training_panel(self):
        """打开训练面板"""
        # 切换到训练标签页
        for i in range(self.tabs.count()):
            if self.tabs.tabText(i) == "训练":
                self.tabs.setCurrentIndex(i)
                break
                
        # 如果有当前项目，自动加载项目数据集
        if hasattr(self, 'training_panel') and self.current_project:
            self.training_panel.use_current_project_dataset()
            
    def open_inference_panel(self):
        """打开推理面板"""
        # 切换到推理标签页
        for i in range(self.tabs.count()):
            if self.tabs.tabText(i) == "推理":
                self.tabs.setCurrentIndex(i)
                break
                
        # 刷新模型列表
        if hasattr(self, 'inference_panel'):
            # 刷新可用模型列表
            self.inference_panel.refresh_model_list()
            
            # 如果有当前项目且包含图像，预先加载一个图像
            if self.current_project:
                images_dir = os.path.join(self.current_project["path"], "images")
                if os.path.exists(images_dir) and os.listdir(images_dir):
                    # 获取第一个图像文件
                    image_files = [f for f in os.listdir(images_dir) 
                                if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
                    if image_files:
                        image_path = os.path.join(images_dir, image_files[0])
                        self.inference_panel.source_path_edit.setText(image_path)

def main():
    """应用程序主入口"""
    app = QApplication(sys.argv)
    
    # 设置应用程序字体
    font = QFont("Microsoft YaHei", 10)
    app.setFont(font)
    
    # 创建资源目录（如果不存在）
    os.makedirs("resources", exist_ok=True)
    
    # 加载样式表
    style_path = os.path.join("resources", "styles.qss")
    
    # 如果样式表文件存在，则加载它
    if os.path.exists(style_path):
        with open(style_path, 'r', encoding='utf-8') as style_file:
            app.setStyleSheet(style_file.read())
    
    # 创建主窗口
    window = MainWindow()
    
    # 显示主窗口
    window.showMaximized()
    
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
